;;; set-randomized-binary-search-tree.ss  --  Jens Axel S?gaard

; Randomized Binary Search Trees

; Reference:
;    "Randomized Binary Search Trees" (1997)
;    Conrado Martinez and Salvador Roura

;;; HISTORY

;  Initial Implemenation

;  22nd jan 2005 [soegaard]
;   - added set?
;  15th jan 2005 [soegaard]
;   - modified remove1 and split/found to handle
;     the case where elm= is finer than elm<, elm>

(module set-randomized-binary-search-tree mzscheme

  (provide randomized-binary-search-tree-set@) 
  
  (require "set-signature.ss"
           "ordered.ss"
           
           (lib "match.ss")
           (lib "list.ss")
           (lib "plt-pretty-big-text.ss" "lang"))
  
  
  ;;; Random
  
  (random-seed 42)  ; Now we can repeat experiments
  
  ; (rnd n) in {0,1,...n}
  (define (rnd n)
    (random (add1 n)))
  
  
  ;;; Randomized Binary Search Tree Set
  
  (define list:foldl foldl)
  
  (define randomized-binary-search-tree-set@
    (unit/sig set^ 
      (import ordered^)

      (define (elm> x y)
        (elm< y x))

      ; DATADEFINITION
      
      ; A RBST is either
      (define empty '())
      ; or
      (define-struct T (l k r s))
      ; where 
      ;   l and r are left and right subtrees.
      ;   k is the key 
      ;   s is the size of the subtree i.e. holds the
      ;     number of elements in the subtree

      (define (set? o)
        (or (empty? o)
            (T? o)))
      
      (define (->sexp s)
        (define -> ->sexp)
        (match s
          [()              '()]
          [($ T l k r s)   `(|T| ,(-> l) ,k ,(-> r) ,(string-append ":" (number->string s)))]))
      
      ; INTERNAL CONSTRUCTORS
      
      ; plain constructor
      (define mt make-T)
      
      ; smart constructor, which takes care of the size invariant
      (define (mt/size l k r)
        (mt l k r (+ (size l) 1 (size r))))
      
      ; CONSTRUCTORS
      
      (define (singleton x)
        (insert empty x))
      
      ; TODO: Use pair-union ?
      (define (set . xs)
        (list:foldl (lambda (x acc) (insert acc x)) empty xs))
      
      ; INSERTION 
      
      (define (insert-all tree xs)
        (list:foldl (lambda (x acc) (insert acc x)) tree xs))
      
      (define (insert tree x)
        (if (member? tree x)   
            tree
            (match tree
              [()            (mt empty x empty 1)]
              [($ T l k r s) (cond
                               [(= (rnd s) s)  (insert-at-root tree x)]
                               [(elm< x k)     (mt (insert l x) k r (add1 s))]
                               [else           (mt l k (insert r x) (add1 s))])])))
      
      
      (define (insert-at-root tree x)
        (let-values ([(t< t>) (split tree x)])
          (mt/size t< x t>)))
      
      ; This split assumes x it not a member of tree.
      (define (split tree x)
        (match tree
          [()             (values () ())]
          [($ T l k r s)  (if (elm< x k)
                              (let-values ([(l< l>) (split l x)])
                                (values l< (mt/size l> k r)))                            
                              (let-values ([(r< r>) (split r x)])
                                (values (mt/size l k r<)  r>)))]))
      
      ; REMOVAL
      
      (define (remove-all tree xs)
        (list:foldl (lambda (x acc) (remove acc x)) tree xs))

      
      ; TODO: PROBLEM: What happens to the size, if x is not in the tree?
      (define (remove tree x)
        (define (rem1 tree)
          (match tree
            [()            empty]
            [($ T l k r s) (cond
                             [(elm< x k)    (mt (rem1 l) k r (sub1 s))]
                             [(elm> x k)    (mt l k (rem1 r) (sub1 s))]
                             [(elm= x k)    (join l r)]
                             [(member? x l) (mt (rem1 l) k r (sub1 s))]
                             [else          (mt l k (rem1 r) (sub1 s))])]))
        (if (member? tree x)
            (rem1 tree)
            tree))

      (define (join t1 t2)
        (let* ([m (size t1)]
               [n (size t2)]
               [total (+ m n)])
          (cond 
            [(= total 0)              empty]
            [(< (rnd (sub1 total)) m) (match t1
                                        [($ T l k r s)  (mt l k (join r t2) total)])]
            [else                     (match t2
                                        [($ T l k r s)  (mt (join t1 l) k r total)])])))
      
      (define (remove-min tree)
        (remove (find-min tree) tree))
      
      (define (remove-max tree)
        (remove (find-max tree) tree))
      
      
      ; SET OPERATIONS
      
      (define (split/found tree x)
        (match tree
          [()             (values () () #f)]
          [($ T l k r s)  (cond
                            [(elm= x k)    (let-values ([(l< l> fl) (split/found l x)]
                                                        [(r< r> fr) (split/found r x)])
                                             (values (join l< l>)
                                                     (join r< r>)
                                                     #t))]
                            [(elm< x k)    (let-values ([(l< l> f) (split/found l x)])
                                             (values l< 
                                                     (mt/size l> k r)
                                                     f))]
                            [(elm> x k)    (let-values ([(r< r> f) (split/found r x)])
                                             (values (mt/size l k r<) 
                                                     r>
                                                     f))]
;                            [(member? x l)  (let-values ([(l< l> f) (split/found l x)])
;                                              (values l< 
;                                                      (mt/size l> k r)
;                                                      f))]
;                            [(member? x r)  (let-values ([(r< r> f) (split/found r x)])
;                                              (values (mt/size l k r<) 
;                                                      r>
;                                                      f))]
;                            [else           (let-values ([(l< l> f1) (split/found l x)]
;                                                         [(r< r> f2) (split/found r x)])
;                                              (values (join l< r<)
;                                                      (insert1 k (join l> r>))
;                                                        (or f1 f2)))]
                            [else         (error "non total order used")]
                            )]))
      
      (define (union t1 t2)
        (let* ([m (size t1)] [n (size t2)] [total (+ m n)])
          (cond 
            [(= total 0)              empty]
            [(< (rnd (sub1 total)) m) (match t1
                                        [($ T l k r s)  (let-values ([(t2< t2> f) (split/found t2 k)])
                                                          (mt/size (union l t2<) k (union r t2>)))])]
            [else                     (match t2
                                        [($ T l k r s)  (let-values ([(t1< t2> f) (split/found t1 k)])
                                                          (mt/size (union l t1<) k (union r t2>)))])])))
      
      
      (define (intersection t1 t2)
        (match t1
          [()             empty]
          [($ T l k r s)  (let-values ([(t2< t2> f) (split/found t2 k)])
                            (let ([il (intersection l t2<)]
                                  [ir (intersection r t2>)])
                              (if f
                                  (mt/size il k ir)
                                  (join il ir))))]))
      
      (define (difference t1 t2)
        (match t1
          [()             empty]
          [($ T l k r s)  (let-values ([(t2< t2> f) (split/found t2 k)])
                            (let ([dl (difference l t2<)]
                                  [dr (difference r t2>)])
                              (if f
                                  (join dl dr)
                                  (mt/size dl k dr))))]))

      
      ; SELECTION
      
      ; select : RBST integer -> element
      ;  return the n'th element of tree with respect to the order 
      ;  (counting from 0)
      (define (select tree n)
        (match tree
          [()             (error "select called on empty tree")]
          [($ T l k r _)  (cond
                            [(= n (size l))  k]
                            [(< n (size l))  (select l n)]
                            [else            (select r (- n (size l) 1))])]))
      
      ; FOLDING
      
      ; fold-ascending : (element alpha -> alpha) alpha tree -> alpha
      (define (fold-ascending f base tree)
        (define (fold- base tree)
          (match tree
            [()            base]
            [($ T l k r _) (fold- (f k (fold- base l))
                                  r)]))
        (fold- base tree))
      
      ; fold-decending : (element alpha -> alpha) alpha tree -> alpha
      (define (fold-decending f base tree)
        (define (fold- base tree)
          (match tree
            [()            base]
            [($ T l k r _) (fold- (f k (fold- base r))
                                  l)]))
        (fold- base tree))
      
      (define fold fold-ascending)
      (define foldl fold)
      
      ; QUERIES
      
      (define empty? null?)
      
      (define (elements tree)
        (fold cons '() tree))
      
      (define (size tree)
        (if (empty? tree)
            0
            (T-s tree)))
      
      (define (find-min tree)
        (match tree
          [()              (error "find-min called on empty tree")]
          [($ T () k _ _)   k]
          [($ T l  _ _ _)  (find-min l)]))
      
      (define (find-max tree)
        (match tree
          [()              (error "find-max called on empty tree")]
          [($ T _ k () _)   k]
          [($ T _ _ r  _)  (find-max r)]))
      
      ; return the key k equal to x otherwise #f
      (define (get tree x)
        (match tree
          [()             #f]
          [($ T l k r _)  (cond
                            [(elm= x k) k]
                            [(elm< x k) (member x l)]
                            [(elm> x k) (member x r)]
                            [else       (or (get l x)
                                            (get r x))])]))
      
      (define (member? tree x)
        (match tree
          [()             #f]
          [($ T l k r _)  (cond
                            [(elm= x k) #t]
                            [(elm< x k) (member? l x)]
                            [(elm> x k) (member? r x)]
                            [else       (or (member? l x)
                                            (member? r x))])]))
      
       
      ; TODO:  Find better algorithm
      (define (subset? s1 s2)
        (fold (lambda (x p?)
                (and p?
                     (member? s2 x)))
              #t s1))
      
      (define (equal? s1 s2)
        (and (subset? s1 s2)
             (subset? s2 s1))))))
  